import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math
import xlrd
import xlwt
import random
from sklearn import ensemble


###########1.读取数据部分##########
# 载入数据并且打乱数据集
# 样本起始行数，结束行数，测试集占总样本集比重,特征数，是否打乱样本集
# #如果Testproportion为0或1就训练集=测试集
def load_data(StartPo, EndPo, TestProportion, FeatureNum, Shuffle, FilePath):
    # 打开excel文件
    # excel路径
    workbook = xlrd.open_workbook(str(FilePath))
    # sheet表
    sheet = workbook.sheet_by_name('Sheet1')
    # 总样本集
    Sample = []
    # 训练集
    train = []
    # 测试集
    test = []
    # 测试集数目
    TestSetSphere = (EndPo - StartPo + 1) * TestProportion
    # 测试集数目
    TestSetSphere = int(TestSetSphere)
    # 获取全部样本集并打乱顺序
    for loadi in range(StartPo - 1, EndPo):
        RowSample = sheet.row_values(loadi)
        Sample.append(RowSample)
    # 是否打乱样本集
    if Shuffle == 1:
        # 如果shuffle=1，打乱样本集
        random.shuffle(Sample)
    # 如果Testproportion为0就训练集=测试集
    if TestProportion == 0 or TestProportion == 1:
        # 变换为array
        TrainSet = np.array(Sample)
        TestSet = np.array(Sample)
    else:
        # 设置训练集
        for loadtraina in Sample[:(EndPo - TestSetSphere)]:
            GetTrainValue = loadtraina
            train.append(GetTrainValue)
        # 设置测试集
        for loadtesta in range(-TestSetSphere - 1, -1):
            GetTestValue = Sample[loadtesta]
            test.append(GetTestValue)
        # 变换样本集
        TrainSet = np.array(train)
        # 变换为array
        TestSet = np.array(test)
    # 分割特征与目标变量
    x1, y1 = TrainSet[:, :FeatureNum], TrainSet[:, -1]
    x2, y2 = TestSet[:, :FeatureNum], TestSet[:, -1]
    return x1, y1, x2, y2


###########2.回归部分##########
def regression_method(model):
    model.fit(x_train, y_train)
    score = model.score(x_test, y_test)
    result = model.predict(x_test)
    # 计算残差平方
    ResidualSquare = (result - y_test) ** 2
    # 计算残差平方和
    RSS = sum(ResidualSquare)
    # 计算均方差
    MSE = np.mean(ResidualSquare)
    # 回归样本个数
    num_regress = len(result)
    print(f'n={num_regress}')
    print(f'R^2={score}')
    print(f'MSE={MSE}')
    print(f'RSS={RSS}')
    ############绘制折线图##########
    plt.figure()
    plt.plot(np.arange(len(result)), y_test, 'go-', label='true value')
    plt.plot(np.arange(len(result)+10), result, 'ro-', label='predict value')
    plt.title('RandomForestRegression R^2: %f' % score)
    plt.legend()  # 将样例显示出来
    plt.show()
    return result


##########3.绘制验证散点图########
def scatter_plot(TureValues, PredictValues):
    # 设置参考的1：1虚线参数
    xxx = [-0.5, 1.5]
    yyy = [-0.5, 1.5]
    # 绘图
    plt.figure()
    # 绘制虚线
    plt.plot(xxx, yyy, c='0', linewidth=1, linestyle=':', marker='.', alpha=0.3)
    # 绘制散点图，横轴是真实值，竖轴是预测值
    plt.scatter(TureValues, PredictValues, s=20, c='r', edgecolors='k', marker='o', alpha=0.8)
    # 设置坐标轴范围
    plt.xlim((0, 1))
    plt.ylim((0, 1))
    plt.title('RandomForestRegressionScatterPlot')
    plt.show()


###########4.预设回归方法##########
####随机森林回归####

# esitimators决策树数量
model_RandomForestRegressor = ensemble.RandomForestRegressor(n_estimators=800)

########5.设置参数与执行部分#############
# 设置数据参数部分
# 行数以excel里为准
x_train, y_train, x_test, y_test = load_data(2, 53, 1, 1, 0, 'D:\___PersonalData\Desktop\data.xls')
# 起始行数2，结束行数121，训练集=测试集，特征数量17,不打乱样本集
# 括号内填上方法，并获取预测值
y_pred = regression_method(model_RandomForestRegressor)
# 生成散点图
scatter_plot(y_test, y_pred)
